# /usr/lib/python3.11/site-packages/langgraph_runtime_postgres/lifespan.pylifespan.py

import asyncio
import signal
from contextlib import asynccontextmanager

import structlog
from langchain_core.runnables.config import RunnableConfig, var_child_runnable_config
from langgraph.constants import CONF
from langgraph_license.validation import (
    check_license_periodically,
    get_license_status,
    plus_features_enabled,
)
from starlette.applications import Starlette

from langgraph_runtime_postgres import checkpoint, database, queue

logger = structlog.stdlib.get_logger(__name__)


@asynccontextmanager
async def lifespan(
    app: Starlette | None = None,
    with_cron_scheduler: bool = True,
    taskset: set[asyncio.Task] | None = None,
):
    # Import at point of use to avoid circular dependencies
    from langgraph_api import (
        __version__,
        config,
        cron_scheduler,
        feature_flags,
        graph,
        http,
        metadata,
        thread_ttl,
    )
    from langgraph_api import asyncio as langgraph_asyncio
    from langgraph_api import (
        store as api_store,
    )
    from langgraph_api.js import ui

    await logger.ainfo(
        f"Starting Postgres runtime with langgraph-api={__version__}",
        version=__version__,
    )

    try:
        current_loop = asyncio.get_running_loop()
        langgraph_asyncio.set_event_loop(current_loop)
    except RuntimeError:
        await logger.aerror("Failed to set loop")

    # Need access to redis client in get_license_status
    await database.start_pool()
#   Update at 2025-07-27
#    if not await get_license_status():
#        raise ValueError(
#            "License verification failed. Please ensure proper configuration:\n"
#            "- For local development, set a valid LANGSMITH_API_KEY for an account with LangGraph Cloud access "
#            "in the environment defined in your langgraph.json file.\n"
#            "- For production, configure the LANGGRAPH_CLOUD_LICENSE_KEY environment variable "
#            "with your LangGraph Cloud license key.\n"
#            "Review your configuration settings and try again. If issues persist, "
#            "contact support for assistance."
#        )
    await http.start_http_client()
    await ui.start_ui_bundler()
    try:
        async with langgraph_asyncio.SimpleTaskGroup(
            cancel=True, taskset=taskset, taskgroup_name="Lifespan"
        ) as tg:
            tg.create_task(metadata.metadata_loop())
            # Update at 2025-07-27
            # tg.create_task(check_license_periodically())
            await api_store.collect_store_from_env()
            store_instance = await api_store.get_store()
            if not api_store.CUSTOM_STORE:
                tg.create_task(store_instance.start_ttl_sweeper())  # type: ignore
            else:
                await logger.ainfo("Using custom store. Skipping store TTL sweeper.")
            tg.create_task(thread_ttl.thread_ttl_sweep_loop())

            if feature_flags.USE_RUNTIME_API:
                from langgraph._internal._constants import CONFIG_KEY_RUNTIME
                from langgraph.runtime import Runtime

                langgraph_config: RunnableConfig = {
                    CONF: {CONFIG_KEY_RUNTIME: Runtime(store=store_instance)}
                }
            else:
                from langgraph.constants import CONFIG_KEY_STORE

                langgraph_config: RunnableConfig = {
                    CONF: {CONFIG_KEY_STORE: store_instance}
                }

            var_child_runnable_config.set(langgraph_config)

            # Keep after the setter above so users can access the store from within the factory function
            await graph.collect_graphs_from_env(True)
            if config.N_JOBS_PER_WORKER > 0:
                tg.create_task(queue_with_signal())
            else:
                await logger.ainfo("N_JOBS_PER_WORKER is 0. Skipping queue.")
            if (
                with_cron_scheduler
                and config.FF_CRONS_ENABLED
                and plus_features_enabled()
            ):
                tg.create_task(cron_scheduler.cron_scheduler())
            tg.create_task(checkpoint.checkpoint_ingestion_loop())
            yield
    finally:
        await api_store.exit_store()
        await ui.stop_ui_bundler()
        await graph.stop_remote_graphs()
        await http.stop_http_client()
        await database.stop_pool()


async def queue_with_signal():
    try:
        await queue.queue()
    except asyncio.CancelledError:
        pass
    except TimeoutError:
        logger.awarning("Timed out waiting for task. Signaling shutdown")
        signal.raise_signal(signal.SIGINT)
    except Exception as exc:
        logger.exception("Queue failed. Signaling shutdown", exc_info=exc)
        signal.raise_signal(signal.SIGINT)

